package com.yangzhe.algorithm.c035;

import java.util.*;

// 插入、删除和获取随机元素O(1)时间且允许有重复数字的结构
public class Code04_InsertDeleteRandomDuplicatesAllowed_LeetCode {

    public static void main(String[] args) {
        RandomizedCollection collection = new RandomizedCollection();
        collection.insert(4);
        collection.insert(3);
        collection.insert(4);
        collection.insert(2);
        collection.insert(4);
        collection.remove(4);
        collection.remove(3);
        collection.remove(4);
        collection.remove(4);
    }

    static class RandomizedCollection {

        /**
         * list 随机访问list[x]速度快，插入删除速度快
         */
        private List<Integer> list;

        /**
         * key 是插入的val， value中的set用来保存list中的index
         * hashmap查找速度快，set查找速度快
         * 由于hashmap不能重复，所以需要再嵌套一层set用来保存多个相同value在list中的index
         */
        private Map<Integer, Set<Integer>> map;

        public RandomizedCollection() {
            list = new ArrayList<>();
            map = new HashMap<>();
        }

        public boolean insert(int val) {
            boolean result;
            Set<Integer> set;
            // 判断是否有数字插入过
            if (map.containsKey(val)) {
                result = false;
                set = map.get(val);
            } else {
                result = true;
                set = new HashSet<>();
            }

            list.add(val);
            set.add(list.size() - 1);
            map.put(val, set);

            return result;
        }

        public boolean remove(int val) {
            if (!map.containsKey(val)) {
                return false;
            }

            Set<Integer> set = map.get(val);
            Integer index = set.iterator().next();
            // 如果要移除的元素不是最后一个，则把最后一个元素覆盖到要移除的位置，相当于删掉了这个元素
            if (index != list.size() - 1) {
                Integer lastValue = list.get(list.size() - 1);
                list.set(index, lastValue);

                Set<Integer> lastSet = map.get(lastValue);
                lastSet.remove(list.size() - 1);
                lastSet.add(index);

                map.put(lastValue, lastSet);

                // 移除最后一个元素
                list.remove(list.size() - 1);

                // 如果最后一个值和要移除的值相等，不能重复从set中删除index，因为之前覆盖操作相当于已经覆盖过了
                if (lastValue != val) {
                    // 把被覆盖的元素index删掉
                    set.remove(index);
                }

                if (set.isEmpty()) {
                    map.remove(val);
                }
            } else {
                // 移除最后一个元素
                list.remove(list.size() - 1);
                // 把被覆盖的元素index删掉
                set.remove(index);

                if (set.isEmpty()) {
                    map.remove(val);
                }
            }


            return true;
        }

        public int getRandom() {
            int index = (int) (list.size() * Math.random());
            return list.get(index);
        }
    }

}
